package com.pig4cloud.pigx.knowledge.support.handler.rag;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.pig4cloud.pigx.knowledge.dto.AiMessageResultDTO;
import com.pig4cloud.pigx.knowledge.dto.ChatMessageDTO;
import com.pig4cloud.pigx.knowledge.entity.AiDatasetEntity;
import com.pig4cloud.pigx.knowledge.service.AiStreamAssistantService;
import com.pig4cloud.pigx.knowledge.support.provider.ModelProvider;
import com.pig4cloud.pigx.knowledge.support.rule.ChatRule;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import reactor.core.publisher.Flux;

import java.util.List;

/**
 * @author lengleng
 * @date 2024/7/4
 */
@Slf4j
@RequiredArgsConstructor
public abstract class AbstractRagChatHandler implements RagChatHandler, ChatRule {

	private final ModelProvider modelProvider;

	@Value("classpath:/prompts/knowledge-system.st")
	private Resource systemResource;

	/**
	 * 结果总结输出
	 * @param dataset 数据
	 * @param chatMessageDTO 聊天消息 DTO
	 * @param resultList 结果列表
	 * @param qdrantResults QDRANT 结果
	 * @return {@link Flux }<{@link AiMessageResultDTO }>
	 */
	public Flux<AiMessageResultDTO> summaryResult(AiDatasetEntity dataset, ChatMessageDTO chatMessageDTO,
			List<String> resultList) {
		// 对结果进行总结
		PromptTemplate userTemplate = new PromptTemplate(systemResource);
		userTemplate.add("contents", CollUtil.join(resultList, StrUtil.CRLF));
		userTemplate.add("userMessage", chatMessageDTO.getContent());
		userTemplate.add("emptyDesc", dataset.getEmptyDesc());

		AiStreamAssistantService streamAssistantService = modelProvider
			.getAiStreamAssistant(chatMessageDTO.getModelName())
			.getValue();
		Flux<String> summaryResult = streamAssistantService.chat(chatMessageDTO.getConversationId(),
				userTemplate.render());
		return summaryResult.map(AiMessageResultDTO::new);
	}

}
